# Copyright (C) 2009  Rob Britton
#
# This file is part of OORegress.
#
# OORegress is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

def parse_params str
  params = []
  buffer = ""

  brackets = 0
  str.split(//).each do |c|
    case c
    when "("
      brackets += 1
      buffer += c
    when ")"
      brackets -= 1
      raise "Unexpected ')'" if brackets < 0
      buffer += c
    when ","
      if brackets == 0
        params << buffer
        buffer = ""
      else
        buffer += ","
      end
    else
      buffer += c
    end
  end

  params << buffer unless buffer == ""

  params.map { |p| p.strip }
end

class ParseNode
  def min_obs
    0
  end
end

class Constant < ParseNode
  def initialize v
    @value = v
  end

  def text
    @value.to_s
  end

  def value(t)
    @value
  end
end

class BuiltIn < ParseNode
  attr_accessor :fun, :params

  def initialize fun, params
    @fun, @params = fun, params

    if fun == "delta"
      raise "Invalid number of parameters to delta" if params.size != 1
    elsif fun == "lag"
      raise "Invalid number of parameters to lag" if params.size > 2 or params.size < 1
      @lag_length = params[1] || 1
    end
  end

  def value(t)
    case @fun
    when "delta"
      raise "Index must be > 0 for a delta variable" if t == 0
      @params.first.value(t) - @params.first.value(t - 1)
    when "lag"
      raise "Index must be greater than the lag length" if t < lag_length(t)
      @params.first.value(t - lag_length(t))
    else
      Math.send(@fun, *@params.map { |p| p.value } )
    end
  end

  def text
    @fun + @params.map { |p| p.text }.join(", ")
  end

  def min_obs
    case @fun
    when "delta"
      1 + params.first.min_obs
    when "lag"
      lag_length + params.first.min_obs
    else
      0
    end
  end

  def lag_length(t = 0)
    @lag_length.is_a?(Expr) ? @lag_length.value(t).to_i : @lag_length
  end
end

class Identifier < ParseNode
  @@dataset = nil

  attr_accessor :name
  
  def initialize name
    @name = name
  end

  def value(t = 0)
    @@dataset[@name][t]
  end

  def self.dataset= s
    @@dataset = s
  end

  def text
    name
  end
end

class Negation < ParseNode
  attr_accessor :param

  def initialize param
    @param = param
  end

  def value(t)
    -param.value(t)
  end

  def text
    "-" + param.text
  end

  def min_obs
    param.min_obs
  end
end

class Expr
  attr_accessor :term, :op, :tail

  def initialize f, o = nil, t = nil
    @term, @op, @tail = f, o, t
  end

  def value(t)
    acc = @term.value(t)
    tail = @tail
    op = @op
    while tail != nil
      if op == :plus
        acc += tail.term.value(t)
      else
        acc -= tail.term.value(t)
      end
      op = tail.op
      tail = tail.tail
    end
    acc
  end

  def value_array(t)
    if @tail
      res = @tail.value_array(t)
      
      res[0] *= (@op == :plus ? 1 : -1)
      [term.value(t)] + res
    else
      [term.value(t)]
    end
  end

  def text
    "(" + @term.text + (@op == nil ? "" :
                        " " + (@op == :plus ? "+" : "-") + " " +
                        @tail.text) + ")"
  end

  def text_array
    [@term.text] + (@op == nil ? [] : @tail.text_array)
  end

  def size
    1 + (@tail == nil ? 0 : @tail.size)
  end

  def min_obs
    if @tail
      t_min = @tail.min_obs
      h_min = @term.min_obs

      (t_min > h_min) ? t_min : h_min
    else
      @term.min_obs
    end
  end
end

class Term
  attr_accessor :factor, :op, :tail

  def initialize f, o = nil, t = nil
    @factor, @op, @tail = f, o, t
  end

  def value(t)
    acc = @factor.value(t)
    tail = @tail
    op = @op
    while tail != nil
      if op == :times
        acc *= tail.factor.value(t)
      else
        acc /= tail.factor.value(t)
      end
      op = tail.op
      tail = tail.tail
    end
    acc
  end

  def text
    @factor.text + (@op == nil ? "" :
                        " " + (@op == :times ? "*" : "/") + " " +
                        @tail.text)
  end

  def min_obs
    if @tail
      t_min = @tail.min_obs
      h_min = @factor.min_obs

      t_min > h_min ? t_min : h_min
    else
      @factor.min_obs
    end
  end
end

class Factor
  attr_accessor :first, :last

  def initialize f, l = nil
    @first, @last = f, l
  end

  def value(t)
    if @last == nil
      @first.value(t)
    else
      @first.value(t) ** @last.value(t)
    end
  end

  def text
    @first.text + (@last == nil ? "" : "^" + @last.text)
  end

  def min_obs
    if @last
      f_min = @first.min_obs
      l_min = @last.min_obs

      (f_min > l_min) ? f_min : l_min
    else
      @first.min_obs
    end
  end
end

class Parser
  BuiltInFunctions = %w[log ln sin cos exp tan delta lag]

  def self.accept char
    raise "Unexpected token '#{@@tokens.first}', expecting '#{char}'" if @@tokens.first != char
    @@tokens.shift
  end

  def self.parseObj
    case @@tokens.first
    when "("
      @@tokens.shift  
      res = parseExpr
      accept ")"
      res
    when /^[a-zA-Z]/  #id or fun
      if BuiltInFunctions.index(@@tokens.first)
        fun = @@tokens.shift
        exprs = []
        accept "("
        exprs << parseExpr

        while @@tokens.first == ","
          accept ","
          exprs << parseExpr
        end

        accept ")"
        BuiltIn.new(fun, exprs)
      else
        Identifier.new(@@tokens.shift)
      end
    when /^\d+/  # a number
      Constant.new @@tokens.shift.to_f
    when /^\-/  # negative something
      accept "-"
      Negation.new(parseObj)
    else
      raise "Unexpected token: '#{@@tokens.first}'"
    end
  end

  def self.parseFactor
    Factor.new parseObj,
      if @@tokens.first == "^"
        accept "^"
        parseObj
      else
        nil
      end
  end

  def self.parseTerm
    head = parseFactor
    op, tail =
        if @@tokens.first == "*"
          @@tokens.shift
          [:times, parseTerm]
        elsif @@tokens.first == "/"
          @@tokens.shift
          [:div, parseTerm]
        else
          nil
        end
    Term.new head, op, tail
  end

  def self.parseExpr
    head = parseTerm
    op, tail =
        if @@tokens.first == "+"
          @@tokens.shift
          [:plus, parseExpr]
        elsif @@tokens.first == "-"
          @@tokens.shift
          [:minus, parseExpr]
        else
          nil
        end
    Expr.new head, op, tail
  end

  def self.parse tokens
    @@tokens = tokens
    lhs = self.parseExpr
    accept "="
    rhs = self.parseExpr

    [lhs, rhs]
  end
end

def parse_eq eq
  tokens = []

  # tokenize
  while eq.length > 0
    eq = eq.sub(/^\s+/, '')
    if eq =~ /^([a-zA-Z][a-zA-Z0-9_]*|\d+|[\+\(\)\-\*\/=\^\,])/
      eq = eq[($1.length)..-1]
      tokens << $1
    else
      raise "Unrecognized token: " + eq[/^\s*([^\s]+)/, 1]
    end
  end
  
  Parser.parse tokens
end
